import os
import pandas as pd
from typing import List, Dict, Any, Optional, Union, Tuple
import logging

# 根据数据库类型导入相应的驱动
try:
    import cx_Oracle as oracle
except ImportError:
    oracle = None

try:
    import pymysql as mysql
except ImportError:
    mysql = None

try:
    import pyodbc as sqlserver
except ImportError:
    sqlserver = None

class DatabaseUtils:
    """数据库操作工具类，支持Oracle、MySQL和SQL Server"""
    
    def __init__(self, db_type: str, host: str, port: int, user: str, password: str, 
                 database: str = None, sid: str = None, dsn: str = None, 
                 charset: str = 'utf8', autocommit: bool = False):
        """
        初始化数据库连接参数
        
        Args:
            db_type: 数据库类型，支持 'oracle', 'mysql', 'sqlserver'
            host: 主机名
            port: 端口号
            user: 用户名
            password: 密码
            database: 数据库名（MySQL/SQL Server）
            sid: Oracle SID
            dsn: SQL Server DSN
            charset: 字符集，默认 'utf8'
            autocommit: 是否自动提交，默认 False
        """
        self.db_type = db_type.lower()
        self.host = host
        self.port = port
        self.user = user
        self.password = password
        self.database = database
        self.sid = sid
        self.dsn = dsn
        self.charset = charset
        self.autocommit = autocommit
        self.connection = None
        
        # 验证数据库类型和驱动
        if self.db_type == 'oracle' and oracle is None:
            raise ImportError("缺少Oracle驱动，请安装 cx_Oracle")
        elif self.db_type == 'mysql' and mysql is None:
            raise ImportError("缺少MySQL驱动，请安装 pymysql")
        elif self.db_type == 'sqlserver' and sqlserver is None:
            raise ImportError("缺少SQL Server驱动，请安装 pyodbc")
    
    def connect(self) -> None:
        """建立数据库连接"""
        try:
            if self.db_type == 'oracle':
                # Oracle 连接
                dsn = oracle.makedsn(self.host, self.port, sid=self.sid)
                self.connection = oracle.connect(
                    user=self.user, 
                    password=self.password, 
                    dsn=dsn
                )
            elif self.db_type == 'mysql':
                # MySQL 连接
                self.connection = mysql.connect(
                    host=self.host,
                    port=self.port,
                    user=self.user,
                    password=self.password,
                    database=self.database,
                    charset=self.charset,
                    autocommit=self.autocommit
                )
            elif self.db_type == 'sqlserver':
                # SQL Server 连接
                if self.dsn:
                    conn_str = f'DSN={self.dsn};UID={self.user};PWD={self.password}'
                else:
                    conn_str = (
                        f'DRIVER={{ODBC Driver 17 for SQL Server}};'
                        f'SERVER={self.host},{self.port};'
                        f'DATABASE={self.database};'
                        f'UID={self.user};PWD={self.password}'
                    )
                self.connection = sqlserver.connect(conn_str)
                
            logging.info(f"成功连接到 {self.db_type.upper()} 数据库")
            
        except Exception as e:
            logging.error(f"数据库连接失败: {e}")
            raise
    
    def close(self) -> None:
        """关闭数据库连接"""
        if self.connection:
            self.connection.close()
            logging.info(f"{self.db_type.upper()} 数据库连接已关闭")
            self.connection = None
    
    def execute(self, sql: str, params: Optional[Union[Tuple, Dict]] = None) -> int:
        """
        执行SQL语句（INSERT、UPDATE、DELETE等）
        
        Args:
            sql: SQL语句
            params: 参数，元组或字典
        
        Returns:
            受影响的行数
        """
        if not self.connection:
            self.connect()
        
        try:
            with self.connection.cursor() as cursor:
                cursor.execute(sql, params or ())
                if not self.autocommit:
                    self.connection.commit()
                return cursor.rowcount
        except Exception as e:
            if not self.autocommit:
                self.connection.rollback()
            logging.error(f"SQL执行失败: {e}\nSQL: {sql}")
            raise
    
    def executemany(self, sql: str, params_list: List[Union[Tuple, Dict]]) -> int:
        """
        批量执行SQL语句
        
        Args:
            sql: SQL语句
            params_list: 参数列表，每个元素是一个元组或字典
        
        Returns:
            受影响的总行数
        """
        if not self.connection:
            self.connect()
        
        try:
            with self.connection.cursor() as cursor:
                cursor.executemany(sql, params_list)
                if not self.autocommit:
                    self.connection.commit()
                return cursor.rowcount
        except Exception as e:
            if not self.autocommit:
                self.connection.rollback()
            logging.error(f"批量SQL执行失败: {e}\nSQL: {sql}")
            raise
    
    def query(self, sql: str, params: Optional[Union[Tuple, Dict]] = None) -> List[Dict[str, Any]]:
        """
        执行查询语句
        
        Args:
            sql: SQL语句
            params: 参数，元组或字典
        
        Returns:
            结果列表，每个元素是一个字典（键为列名，值为列值）
        """
        if not self.connection:
            self.connect()
        
        try:
            with self.connection.cursor() as cursor:
                cursor.execute(sql, params or ())
                
                # 获取列名
                columns = [col[0] for col in cursor.description]
                
                # 转换为字典列表
                return [dict(zip(columns, row)) for row in cursor.fetchall()]
        except Exception as e:
            logging.error(f"查询失败: {e}\nSQL: {sql}")
            raise
    
    def query_one(self, sql: str, params: Optional[Union[Tuple, Dict]] = None) -> Optional[Dict[str, Any]]:
        """
        执行查询并返回单条记录
        
        Args:
            sql: SQL语句
            params: 参数，元组或字典
        
        Returns:
            单条记录（字典）或 None
        """
        if not self.connection:
            self.connect()
        
        try:
            with self.connection.cursor() as cursor:
                cursor.execute(sql, params or ())
                
                # 获取列名
                columns = [col[0] for col in cursor.description]
                
                # 获取单条记录
                row = cursor.fetchone()
                return dict(zip(columns, row)) if row else None
        except Exception as e:
            logging.error(f"查询失败: {e}\nSQL: {sql}")
            raise
    
    def query_df(self, sql: str, params: Optional[Union[Tuple, Dict]] = None) -> pd.DataFrame:
        """
        执行查询并返回DataFrame
        
        Args:
            sql: SQL语句
            params: 参数，元组或字典
        
        Returns:
            Pandas DataFrame
        """
        if not self.connection:
            self.connect()
        
        try:
            return pd.read_sql(sql, self.connection, params=params)
        except Exception as e:
            logging.error(f"查询失败: {e}\nSQL: {sql}")
            raise
    
    def call_procedure(self, proc_name: str, params: Optional[List[Any]] = None) -> None:
        """
        调用存储过程
        
        Args:
            proc_name: 存储过程名称
            params: 参数列表
        """
        if not self.connection:
            self.connect()
        
        try:
            with self.connection.cursor() as cursor:
                cursor.callproc(proc_name, params or [])
                if not self.autocommit:
                    self.connection.commit()
        except Exception as e:
            if not self.autocommit:
                self.connection.rollback()
            logging.error(f"存储过程调用失败: {e}\n过程名: {proc_name}")
            raise
    
    def __enter__(self):
        """支持 with 语句"""
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """支持 with 语句"""
        self.close()
        return False  # 不抑制异常


# 使用示例
if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
    
    # MySQL 示例
    try:
        with DatabaseUtils(
            db_type='mysql',
            host='localhost',
            port=3306,
            user='root',
            password='password',
            database='test_db'
        ) as db:
            # 创建表
            db.execute('''
            CREATE TABLE IF NOT EXISTS users (
                id INT PRIMARY KEY AUTO_INCREMENT,
                name VARCHAR(50) NOT NULL,
                age INT
            )
            ''')
            
            # 插入数据
            db.execute("INSERT INTO users (name, age) VALUES (%s, %s)", ('Alice', 25))
            
            # 批量插入
            users = [('Bob', 30), ('Charlie', 35)]
            db.executemany("INSERT INTO users (name, age) VALUES (%s, %s)", users)
            
            # 查询数据
            results = db.query("SELECT * FROM users WHERE age > %s", (30,))
            print("年龄大于30的用户:")
            for row in results:
                print(f"ID: {row['id']}, 姓名: {row['name']}, 年龄: {row['age']}")
            
            # 使用DataFrame查询
            df = db.query_df("SELECT * FROM users")
            print("\n所有用户数据:")
            print(df)
            
            # 调用存储过程（如果有）
            # db.call_procedure("sp_test", [param1, param2])
    
    except Exception as e:
        print(f"操作失败: {e}")
    
    # Oracle 示例（需要正确配置Oracle环境）
    # try:
    #     with DatabaseUtils(
    #         db_type='oracle',
    #         host='localhost',
    #         port=1521,
    #         user='system',
    #         password='password',
    #         sid='ORCL'
    #     ) as db:
    #         # 执行Oracle操作
    #         pass
    # except Exception as e:
    #     print(f"Oracle操作失败: {e}")
    
    # SQL Server 示例（需要正确配置ODBC驱动）
    # try:
    #     with DatabaseUtils(
    #         db_type='sqlserver',
    #         host='localhost',
    #         port=1433,
    #         user='sa',
    #         password='password',
    #         database='test_db'
    #     ) as db:
    #         # 执行SQL Server操作
    #         pass
    # except Exception as e:
    #     print(f"SQL Server操作失败: {e}")    